package com.aqie.medium.tree;

import com.aqie.easy.structure.TreeNode;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * 1110 删点成林
 */
public class DelNodes {
    /**
     * 1, DFS  9ms
     */
    List<TreeNode> list = new ArrayList<>();

    public  List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        Set<Integer> set = new HashSet<>();
        for (int de : to_delete)
            set.add(de);
        if (!set.contains(root.val))
            list.add(root);
        dfs(set, root);

        for (int i = list.size() - 1; i >= 0; i--) {
            if (set.contains(list.get(i).val)){
                list.remove(list.get(i));
                continue;
            }
            deleteNode(set, list.get(i));
        }
        return list;
    }

    private void deleteNode(Set<Integer> set, TreeNode root) {
        if (root == null)
            return;
        if (root.left != null && set.contains(root.left.val))
            root.left = null;
        else
            deleteNode(set, root.left);
        if (root.right != null && set.contains(root.right.val))
            root.right = null;
        else
            deleteNode(set, root.right);
    }

    public  void dfs(Set<Integer> set, TreeNode root) {
        if (root == null)
            return;
        if (set.contains(root.val)) {
            if (isExist(set, root.left))
                dfs(set, root.left);
            if (root.left != null)
                list.add(root.left);
            if (isExist(set, root.right))
                dfs(set, root.right);
            if (root.right != null)
                list.add(root.right);
        } else {
            dfs(set, root.left);
            dfs(set, root.right);
        }
    }

    private  boolean isExist(Set<Integer> set, TreeNode root) {
        if (root == null)
            return false;
        if (set.contains(root.val))
            return true;
        return isExist(set, root.left) || isExist(set, root.right);
    }


    /**
     * 2. 后序遍历 6ms
     */
    Set<Integer> set;
    List<TreeNode> res;
    public List<TreeNode> delNodes2(TreeNode root, int[] to_delete) {
        set=new HashSet();
        res=new ArrayList();
        for(int t:to_delete)
            set.add(t);
        root=helper(root);
        if(root!=null)res.add(root);
        return res;
    }
    private TreeNode helper(TreeNode root){
        if(root==null)return null;
        if(root.left==null&&root.right==null)
            return set.contains(root.val)?null:root;
        TreeNode left=helper(root.left);
        TreeNode right=helper(root.right);
        if(set.contains(root.val)){
            if(left!=null)res.add(left);
            if(right!=null)res.add(right);
            return null;
        }
        root.left=left;
        root.right=right;
        return root;
    }
}
